import torch
import demo

a = torch.ones(size=(3, 1))
b = torch.ones(size=(3, 1))

print(demo.custom_add(a, b))